"""DHF模型汇流计算模块

实现大伙房模型的汇流计算，包括经验单位线汇流和前期影响净雨的计算。
"""

import numpy as np
from typing import List, Dict, Any, Tuple
from scipy import integrate
from .parameters import DHFParameters


class Routing:
    """DHF模型汇流计算类"""
    
    def __init__(self, parameters: DHFParameters, rainfall_center: str = '中游'):
        """
        初始化汇流计算模块
        
        Args:
            parameters: DHF模型参数对象
            rainfall_center: 降雨中心位置 ('上游', '中游', '下游')
        """
        self.params = parameters
        self.routing_params = parameters.routing_params
        self.rainfall_center = rainfall_center
        self.center_params = parameters.get_routing_params_by_center(rainfall_center)
        
        # 初始化前期影响净雨量
        self.y_a = 0.0
        
        # 存储历史径流数据用于汇流计算
        self.runoff_history = []
        self.time_history = []
    
    def calculate_routing_time(self, y_a: float, y_tr: float) -> float:
        """
        计算汇流时间 T_M
        
        Args:
            y_a: 前期影响净雨量 (mm)
            y_tr: 当前时段净雨量 (mm)
            
        Returns:
            汇流时间 T_M (小时)
        """
        L_B = self.routing_params['L_B']
        K0 = self.routing_params['K0']
        
        # 防止除零错误
        total_runoff = max(0.001, y_a + y_tr)
        
        T_M = L_B * (total_runoff ** (-K0))
        
        return max(1.0, T_M)  # 确保T_M至少为1小时
    
    def update_antecedent_runoff(self, y_tr: float):
        """
        更新前期影响净雨量
        
        Args:
            y_tr: 当前时段净雨量 (mm)
        """
        K = self.routing_params['Ka']
        
        # 更新前期影响净雨量
        self.y_a = K * (self.y_a + y_tr)
    
    def calculate_K3_and_AA(self, T_M: float, t0: float = None) -> Tuple[float, float]:
        """
        计算汇流参数 K3 和 AA
        
        Args:
            T_M: 汇流时间 (小时)
            t0: 峰现时间 (小时)，如果为None则取T_M/3
            
        Returns:
            (K3, AA) 参数元组
        """
        n = self.routing_params['N']
        CC = self.center_params['Cc']
        DD = self.center_params['Dd']
        
        if t0 is None:
            t0 = T_M / 3.0  # 默认峰现时间
        
        # 计算AA (根据文档中的公式)
        if t0 > 0 and T_M > 0:
            pi_ratio = np.pi * t0 / (n * T_M)
            if pi_ratio < np.pi/2:  # 确保tan函数有意义
                AA = (CC / DD) * (pi_ratio ** (DD - 1)) * np.tan(pi_ratio)
            else:
                AA = 1.0
        else:
            AA = 1.0
        
        # 计算K3 (通过数值积分)
        def integrand(t):
            if t <= 0 or t >= n * T_M:
                return 0
            pi_t_ratio = np.pi * t / (n * T_M)
            return np.exp(-AA * (pi_t_ratio ** n)) * (np.sin(pi_t_ratio) ** CC)
        
        try:
            integral_result, _ = integrate.quad(integrand, 0, n * T_M)
            if integral_result > 0:
                K3 = (n * T_M) / integral_result
            else:
                K3 = 1.0
        except:
            K3 = 1.0
        
        return K3, AA
    
    def calculate_unit_hydrograph(self, T_M: float, dt: float = 1.0, 
                                 duration: float = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        计算单位线 (公式34)
        
        Args:
            T_M: 汇流时间 (小时)
            dt: 时间步长 (小时)
            duration: 单位线持续时间 (小时)，如果为None则自动计算
            
        Returns:
            (time_array, unit_hydrograph) 时间数组和单位线数组
        """
        n = self.routing_params['N']
        CC = self.center_params['Cc']
        
        if duration is None:
            duration = n * T_M
        
        # 创建时间数组
        time_array = np.arange(0, duration + dt, dt)
        
        # 计算K3和AA
        K3, AA = self.calculate_K3_and_AA(T_M)
        
        # 计算单位线 (公式34)
        unit_hydrograph = np.zeros_like(time_array)
        
        for i, t in enumerate(time_array):
            if t > 0 and t < n * T_M:
                pi_t_ratio = np.pi * t / (n * T_M)
                unit_hydrograph[i] = (K3 / (n * T_M)) * \
                                   np.exp(-AA * (pi_t_ratio ** n)) * \
                                   (np.sin(pi_t_ratio) ** CC)
        
        # 归一化单位线
        total_area = np.trapz(unit_hydrograph, time_array)
        if total_area > 0:
            unit_hydrograph = unit_hydrograph / total_area
        
        return time_array, unit_hydrograph
    
    def convolve_runoff(self, runoff_series: List[float], time_step: float = 1.0) -> np.ndarray:
        """
        对径流序列进行汇流卷积计算
        
        Args:
            runoff_series: 径流序列 (mm)
            time_step: 时间步长 (小时)
            
        Returns:
            汇流后的流量序列 (m³/s)
        """
        if not runoff_series:
            return np.array([])
        
        # 假设流域面积 (km²)，这里使用一个典型值，实际应用中需要输入
        watershed_area = 1000.0  # km²
        
        # 初始化输出序列
        output_length = len(runoff_series) * 3  # 延长输出序列以包含完整的汇流过程
        flow_series = np.zeros(output_length)
        
        # 对每个时段的径流进行汇流计算
        for i, runoff in enumerate(runoff_series):
            if runoff <= 0:
                continue
            
            # 更新前期影响净雨量
            self.update_antecedent_runoff(runoff)
            
            # 计算汇流时间
            T_M = self.calculate_routing_time(self.y_a, runoff)
            
            # 计算单位线
            time_array, unit_hydrograph = self.calculate_unit_hydrograph(T_M, time_step)
            
            # 进行卷积计算
            for j, uh_value in enumerate(unit_hydrograph):
                output_index = i + j
                if output_index < output_length:
                    # 将径流深度转换为流量 (mm -> m³/s)
                    flow_contribution = runoff * uh_value * watershed_area * 1000 / (3600 * time_step)
                    flow_series[output_index] += flow_contribution
        
        return flow_series[:len(runoff_series) + int(max(50, len(runoff_series)))]
    
    def calculate_routing_for_components(self, runoff_components: Dict[str, List[float]], 
                                       time_step: float = 1.0) -> Dict[str, np.ndarray]:
        """
        对不同径流分量分别进行汇流计算
        
        Args:
            runoff_components: 径流分量字典，包含各分量的时间序列
            time_step: 时间步长 (小时)
            
        Returns:
            各分量汇流后的流量序列字典
        """
        routed_components = {}
        
        # 地面壤中流汇流 (使用地面径流参数)
        if 'interflow' in runoff_components:
            routed_components['interflow'] = self.convolve_runoff(
                runoff_components['interflow'], time_step
            )
        
        # 地下径流汇流 (使用地下径流参数，汇流时间更长)
        if 'groundwater_runoff' in runoff_components:
            # 保存当前参数
            original_n = self.routing_params['N']
            
            # 地下径流汇流时间更长
            self.routing_params['N'] = original_n * 2.0
            
            routed_components['groundwater_runoff'] = self.convolve_runoff(
                runoff_components['groundwater_runoff'], time_step
            )
            
            # 恢复原参数
            self.routing_params['N'] = original_n
        
        # 不透水面积径流 (直接流出，不进行汇流)
        if 'impervious_runoff' in runoff_components:
            watershed_area = 1000.0  # km²
            impervious_flow = np.array(runoff_components['impervious_runoff'])
            # 转换为流量 (mm -> m³/s)
            routed_components['impervious_runoff'] = impervious_flow * watershed_area * 1000 / 3600
        
        return routed_components
    
    def calculate_total_flow(self, routed_components: Dict[str, np.ndarray]) -> np.ndarray:
        """
        计算总流量
        
        Args:
            routed_components: 各分量汇流后的流量序列字典
            
        Returns:
            总流量序列 (m³/s)
        """
        # 找到最长的序列长度
        max_length = max(len(flow) for flow in routed_components.values())
        
        # 初始化总流量序列
        total_flow = np.zeros(max_length)
        
        # 累加各分量
        for component_name, flow_series in routed_components.items():
            # 扩展较短的序列
            extended_flow = np.zeros(max_length)
            extended_flow[:len(flow_series)] = flow_series
            total_flow += extended_flow
        
        return total_flow
    
    def get_peak_flow_info(self, flow_series: np.ndarray, time_step: float = 1.0) -> Dict[str, float]:
        """
        获取洪峰流量信息
        
        Args:
            flow_series: 流量序列 (m³/s)
            time_step: 时间步长 (小时)
            
        Returns:
            洪峰信息字典
        """
        if len(flow_series) == 0:
            return {'peak_flow': 0, 'peak_time': 0, 'peak_index': 0}
        
        peak_index = np.argmax(flow_series)
        peak_flow = flow_series[peak_index]
        peak_time = peak_index * time_step
        
        return {
            'peak_flow': peak_flow,      # 洪峰流量 (m³/s)
            'peak_time': peak_time,      # 峰现时间 (小时)
            'peak_index': peak_index     # 峰值索引
        }
    
    def calculate_runoff_volume(self, flow_series: np.ndarray, time_step: float = 1.0) -> float:
        """
        计算径流总量
        
        Args:
            flow_series: 流量序列 (m³/s)
            time_step: 时间步长 (小时)
            
        Returns:
            径流总量 (万m³)
        """
        if len(flow_series) == 0:
            return 0.0
        
        # 计算径流总量 (m³/s * 小时 * 3600s/小时 / 10000 = 万m³)
        total_volume = np.trapz(flow_series, dx=time_step) * 3600 / 10000
        
        return total_volume
    
    def reset_state(self):
        """
        重置汇流计算状态
        """
        self.y_a = 0.0
        self.runoff_history = []
        self.time_history = []
    
    def get_routing_parameters(self) -> Dict[str, Any]:
        """
        获取当前汇流参数
        
        Returns:
            汇流参数字典
        """
        return {
            'routing_params': self.routing_params,
            'center_params': self.center_params,
            'rainfall_center': self.rainfall_center,
            'antecedent_runoff': self.y_a
        }
    
    def set_watershed_area(self, area: float):
        """
        设置流域面积
        
        Args:
            area: 流域面积 (km²)
        """
        self.watershed_area = area
    
    def __str__(self) -> str:
        """汇流模块信息的字符串表示"""
        return f"DHF汇流模块 (降雨中心: {self.rainfall_center}, 前期影响净雨: {self.y_a:.2f}mm)"
    
    def __repr__(self) -> str:
        return self.__str__()